#!/usr/bin/env python
# -*- coding: utf-8 -*-

import rospy
import socket
import json
import threading
import math
from sensor_msgs.msg import NavSatFix
from geometry_msgs.msg import PoseStamped
from std_msgs.msg import Header
from mavros_msgs.msg import State
from mavros_msgs.srv import CommandBool, SetMode, ParamSet # 用于解锁无人机、设置无人机模式、设置参数
from mavros_msgs.msg import PositionTarget

# 全局变量存储处理后的目标GPS数据
target_gps_data = NavSatFix()
target_gps_data.header = Header()
target_gps_data.header.stamp = 0 # 初始时间戳为0，时间戳和本系统时间相差过大的数据手动过滤掉
target_gps_data_lock = threading.Lock()  # 用于线程安全的锁

current_state = State() # 全局变量存储无人机当前状态
global_pos = NavSatFix() # 全局变量存储实际GPS数据
target_global_pos = NavSatFix() # 全局变量存储目标GPS数据

local_pos = PoseStamped() # 全局变量存储无人机本地当前位置
last_local_pos = PoseStamped() # 全局变量存储上一次本地无人机位置
target_local_pos = PositionTarget() # 全局变量存储本地目标位置
# 设置目标位置的相关参数
target_local_pos.header.frame_id = "base_footprint"  # 设置坐标系的标识符
target_local_pos.coordinate_frame = PositionTarget.FRAME_LOCAL_NED  # 使用局部NED坐标系
# 设置掩码，忽略加速度和偏航角速率，仅控制位置、速度和偏航角
target_local_pos.type_mask = (
    PositionTarget.IGNORE_AFX |
    PositionTarget.IGNORE_AFY |
    PositionTarget.IGNORE_AFZ |
    PositionTarget.IGNORE_YAW_RATE
)

def wgs84_displacement(lat_deg, lon_deg, alt, de, dn, du):
    """
    @brief 计算 WGS84 椭球模型下东/北/上方向位移后的新坐标
    :param lat_deg: 原纬度（十进制度，北纬为正）
    :param lon_deg: 原经度（十进制度，东经为正）
    :param alt:       原海拔高（米）
    :param de:      东向位移（米）
    :param dn:      北向位移（米）
    :param du:      上向位移（米）
    :return: (新纬度, 新经度, 新海拔)
    """
    # WGS84 参数
    a = 6378137.0          # 长半轴 [m]
    e2 = 6.69437999014e-3  # 第一偏心率平方

    # 将纬度转换为弧度
    lat_rad = math.radians(lat_deg)
    sin_lat = math.sin(lat_rad)

    # 计算卯酉圈曲率半径 Rn 和子午圈曲率半径 Rm
    Rn = a / math.sqrt(1 - e2 * sin_lat**2)
    Rm = a * (1 - e2) / (1 - e2 * sin_lat**2)**1.5

    # 计算纬度变化（北向位移）
    delta_lat = dn / Rm
    new_lat_deg = lat_deg + math.degrees(delta_lat)

    # 计算经度变化（东向位移）
    delta_lon = de / (Rn * math.cos(lat_rad))
    new_lon_deg = lon_deg + math.degrees(delta_lon)

    # 更新高度
    new_alt = alt + du

    return (new_lat_deg, new_lon_deg, new_alt)

def relative_position_direct(latA_deg, lonA_deg, altA, latB_deg, lonB_deg, altB):
    """
    @brief 计算两GPS点相对位置
    :param latA_deg: 点A纬度（十进制度，北纬为正）
    :param lonA_deg: 点A经度（十进制度，东经为正）
    :param altA: 点A海拔（米）
    :param latB_deg: 点B纬度（十进制度，北纬为正）
    :param lonB_deg: 点B经度（十进制度，东经为正）
    :param altB: 点B海拔（米）
    :return: (东向位移（米），北向位移（米），上向位移（米）)
    """
    # WGS84 参数
    a = 6378137.0          # 长半轴 [m]
    e2 = 6.69437999014e-3  # 第一偏心率平方

    # 转换为弧度
    latA = math.radians(latA_deg)
    lonA = math.radians(lonA_deg)
    latB = math.radians(latB_deg)
    lonB = math.radians(lonB_deg)

    # 计算曲率半径
    RnA = a / math.sqrt(1 - e2 * math.sin(latA)**2)
    RmA = a * (1 - e2) / (1 - e2 * math.sin(latA)**2)**1.5

    # 南北距离（纬度差 × Rm）
    dn = (latB - latA) * RmA

    # 东西距离（经度差 × Rn * cos(latA)）
    de = (lonB - lonA) * RnA * math.cos(latA)

    # 垂直距离
    du = altB - altA

    return (de, dn, du)

# 回调函数，用于更新当前状态
def state_callback(state):
    """
    @brief 状态回调函数
    :param state: 无人机当前状态
    :return: 无返回值
    """
    global current_state
    current_state = state

def local_pose_callback(msg):
    """
    @brief 本地位置回调函数,获取无人机的本地位置
    :param msg: 本地位置消息
    :return: 无返回值
    """
    global local_pos
    global last_local_pos
    last_local_pos = local_pos
    local_pos = msg
    rospy.loginfo("local_pose_callback local_pos.pose.position.x: %.6f, local_pos.pose.position.y: %.6f, local_pos.pose.position.z: %.6f", local_pos.pose.position.x, local_pos.pose.position.y, local_pos.pose.position.z)
    pass

# 接收实际GPS数据并处理
def global_pos_Callback(msg):
    """
    @brief GPS位置回调函数,获取实际GPS数据,并根据目标GPS数据计算本地目标位置
    :param msg: GPS位置消息
    :return: 无返回值
    """
    global global_pos
    global target_local_pos
    global target_gps_data
    global target_global_pos
    global local_pos
    global_pos = msg
    new_d = (0,0,0)
    if target_gps_data.header.stamp != 0:  # 初始时间戳为0，时间戳和本系统时间相差过大的数据手动过滤掉
        with target_gps_data_lock:
            target_global_pos.latitude= target_gps_data.latitude
            target_global_pos.longitude= target_gps_data.longitude
            target_global_pos.altitude= target_gps_data.altitude

        new_d = relative_position_direct(global_pos.latitude, global_pos.longitude, global_pos.altitude, target_global_pos.latitude, target_global_pos.longitude, target_global_pos.altitude)
    target_local_pos.position.x = local_pos.pose.position.x + new_d[1]
    target_local_pos.position.y = local_pos.pose.position.y + new_d[0]
    target_local_pos.position.z = local_pos.pose.position.z + new_d[2]
    # rospy.loginfo(f"target_global_pos.latitude: {target_global_pos.latitude}, target_global_pos.longitude: {target_global_pos.longitude}, target_global_pos.altitude: {target_global_pos.altitude}")
    rospy.logwarn(f"target_local_pos.pose.position.x: {target_local_pos.position.x}, target_local_pos.pose.position.y: {target_local_pos.position.y}, target_local_pos.pose.position.z: {target_local_pos.position.z}")
    rospy.loginfo(f"new_d[1]:{new_d[1]},new_d[0]{new_d[0]},new_d[2]{new_d[2]}")
    rospy.loginfo("global_pos_Callback:Latitude: %f, Longitude: %f, Altitude: %f", msg.latitude, msg.longitude, msg.altitude)



def parse_gps_data(raw_data):
    """
    @brief 解析socket接收到的GPS数据JSON数据
    return: (纬度, 经度, 海拔, 时间戳)
    """
    try:
        # 解析JSON数据
        gps_data = json.loads(raw_data)
        
        # 提取经度、纬度、海拔和时间戳
        latitude = gps_data.get("latitude")
        longitude = gps_data.get("longitude")
        altitude = gps_data.get("altitude")
        timestamp = gps_data.get("timestamp")
        
        if latitude is not None and longitude is not None and altitude is not None:
            return latitude, longitude, altitude, timestamp
        else:
            rospy.logwarn("Invalid GPS data: missing required fields")
            return None
    except json.JSONDecodeError:
        rospy.logwarn("Failed to parse GPS data as JSON")
        return None

def udp_receiver():
    """
    @brief UDP接收线程函数,用于接收基点GPS数据,并计算目标GPS数据，写入全局变量target_gps_data
    :return: 无返回值
    """
    global target_gps_data
    # 创建UDP Socket客户端
    client_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    client_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    server_address = ('0.0.0.0', 12345)  # 绑定到所有网卡，端口12345
    client_socket.bind(server_address)
    rospy.loginfo("UDP receiver thread started, waiting for data...")

    while not rospy.is_shutdown():
        # 接收数据
        raw_data, addr = client_socket.recvfrom(1024)  # 接收数据
        raw_data = raw_data.decode('utf-8')
        rospy.loginfo(f"Received data from {addr}: {raw_data}")
        
        # 解析GPS数据
        gps_data = parse_gps_data(raw_data)
        if gps_data:
            latitude, longitude, altitude, timestamp = gps_data
            
            # 时间相差不大就认为数据有效，处理数据
            if  (rospy.Time.from_sec(timestamp) - rospy.Time.now() < rospy.Duration(3)) and (rospy.Time.from_sec(timestamp) - rospy.Time.now() > rospy.Duration(-3)):  # 时间戳和本系统时间相差过大的数据手动过滤掉
                new_position = wgs84_displacement(latitude, longitude, altitude, 5, 0, 0) # 东向位移5米，北向位移0米，上向位移0米（位于基点东边5米的GPS位置），作为无人机的目标位置
                rospy.loginfo(f"Processed GPS data: {new_position}")
                # 创建NavSatFix消息
                gps_msg = NavSatFix()
                gps_msg.header = Header()
                gps_msg.header.stamp = rospy.Time.from_sec(timestamp)  # 将时间戳转换为ROS时间
                gps_msg.header.frame_id = "target_gps"
                gps_msg.latitude = new_position[0]
                gps_msg.longitude = new_position[1]
                gps_msg.altitude = new_position[2]
                
                # 将处理后的数据写入全局变量
                with target_gps_data_lock:
                    target_gps_data = gps_msg

    # 关闭Socket连接
    client_socket.close()

def gps_client():
    """
    @brief GPS客户端函数,主逻辑的处理,用于启动UDP接收线程、订阅无人机状态、订阅无人机本地位置、订阅实际GPS数据、发布本地目标位置
    :return: 无返回值
    """
    global target_local_pos
    global target_gps_data
    rospy.init_node('gps_receiver', anonymous=True)

    # 启动UDP接收线程
    udp_thread = threading.Thread(target=udp_receiver)
    udp_thread.daemon = True  # 设置为守护线程，主线程退出时自动结束
    udp_thread.start()

    state_sub = rospy.Subscriber('mavros/state', State, state_callback)
    local_pos_sub = rospy.Subscriber('/mavros/local_position/pose', PoseStamped, local_pose_callback)
    global_pos_sub = rospy.Subscriber('/mavros/global_position/global', NavSatFix, global_pos_Callback)

    local_pos_pub = rospy.Publisher('/mavros/setpoint_raw/local', PositionTarget, queue_size=10)

    # 创建服务客户端，用于解锁无人机
    arming_client = rospy.ServiceProxy("/mavros/cmd/arming", CommandBool)
    # 创建服务客户端，用于切换飞行模式
    set_mode_client = rospy.ServiceProxy('mavros/set_mode', SetMode)
    param_set_client = rospy.ServiceProxy('mavros/param/set', ParamSet)

    rate20 = rospy.Rate(20)  # 20 Hz
    rate5 = rospy.Rate(5)  # 5 Hz
    # 等待飞控和MAVROS建立连接
    while not rospy.is_shutdown() and not current_state.connected:
        rate5.sleep()
    rospy.loginfo("已连接到无人机")
    target_local_pos.position.x = 0     # 初始位置为当前位置
    target_local_pos.position.y = 0
    target_local_pos.position.z = 0
    target_local_pos.yaw = 0

    # 在 OFFBOARD 模式下需要先发送一些目标点
    for i in range(50):
        local_pos_pub.publish(target_local_pos)
        rate20.sleep()
    

    # 切换到 OFFBOARD 模式
    if current_state.mode != "OFFBOARD":
        response = set_mode_client(base_mode=0, custom_mode="OFFBOARD")
        rospy.logwarn(f"response:{response}")
        if response.mode_sent:
            rospy.loginfo("成功切换到 OFFBOARD 模式")
        else:
            rospy.logwarn("切换到 OFFBOARD 模式失败")
    else:
        rospy.loginfo("无人机过去已处于 OFFBOARD 模式")
    # 解锁无人机
    if not current_state.armed:
        response = arming_client(True)
        if response.success:
            rospy.loginfo("无人机解锁成功")
        else:
            rospy.logwarn("无人机解锁失败")
    else:
        rospy.loginfo("无人机过去已解锁")



    while not rospy.is_shutdown():
        rospy.logwarn(f"target_gps_data.header.stamp:{target_gps_data.header.stamp}")
        # 如果存在处理后的数据，发布到ROS话题
        if target_local_pos is not None:
            local_pos_pub.publish(target_local_pos)
            rospy.loginfo(f"Published local data: {target_local_pos}")

        # 短暂休眠，避免占用过多CPU资源
        rate20.sleep()

if __name__ == '__main__':
    try:
        gps_client()
    except rospy.ROSInterruptException:
        pass
